# -*- coding: utf-8 -*-

from genericpath import exists
import sys
import os
import cv2 as cv
import pandas as pd
import numpy as np
import shutil
import random

def create_split_dir(data_root, labels):    

    dir_path = os.path.join(data_root, 'data')
    try:
        shutil.rmtree(dir_path)
        print ("path removed", dir_path)
    except OSError as e:
        print("Error: %s : %s" % (dir_path, e.strerror))

    os.mkdir(os.path.join(data_root,'data'))
    os.mkdir(os.path.join(data_root,'data','train'))
    os.mkdir(os.path.join(data_root,'data','test'))
    os.mkdir(os.path.join(data_root,'data','val'))
    for i in range(len(labels.Class_name)):

        path1 = os.path.join(data_root,'data','train',labels.Class_name[i])
        path2 = os.path.join(data_root,'data','test',labels.Class_name[i])
        path3 = os.path.join(data_root,'data','val',labels.Class_name[i])
        # path4 = os.path.join(path,'Raw_Data',labels.Class_name[i])

        path_list = [path1,path2,path3]#,path4
        for dirs in path_list:
            print (dirs)
            if not os.path.exists(dirs):
                os.mkdir(dirs)
            else:
                print('\nDirectory already exists, please delete it!\n')
                sys.exit(0)

def create_label(raw_path, label_path):
    data = os.listdir(raw_path)
    data.sort()

    n = 0
    cls_list = []
    with open(label_path, "w") as file:
        # file.write("Class_id Class_name\r\n")
        for name in data:
            cls_name = name.split('_')[-3]
            if cls_name not in cls_list:
                file.write("{} {}\r\n".format(n, cls_name))
                cls_list.append(cls_name)
                n=n+1

def prepare_clips(data_root, raw_data_path, split, min_length):
    """ 
        Before run this function, you should put your video at {raw_data_path}
        The dumped frames will be save at {data_root}/data. 
    """
    
    video_names = os.listdir(raw_data_path)
    random.shuffle(video_names)

    # calculate number of cls 
    num_train = int(len(video_names)*split[0])
    num_val = int(len(video_names)*split[1])
    num_test = int(len(video_names)-num_train-num_val) 
    # mkdir data_root/data
    if not os.path.exists(os.path.join(data_root, 'data')):
        os.mkdir(os.path.join(data_root, 'data'))
    # mkdir data_root/data/train
    if not os.path.exists(os.path.join(data_root, 'data', 'train')):
        os.mkdir(os.path.join(data_root, 'data', 'train'))
    # mkdir data_root/data/test
    if not os.path.exists(os.path.join(data_root, 'data','test')):
        os.mkdir(os.path.join(data_root, 'data', 'test'))
    # mkdir data_root/data/val
    if not os.path.exists(os.path.join(data_root, 'data', 'val')):
        os.mkdir(os.path.join(data_root, 'data', 'val'))

    cnt = 0
    for video_name in video_names:
        video_path = os.path.join(raw_data_path, video_name)
        video_name = video_name.split('.')[0]

        cap = cv.VideoCapture(video_path)

        num_frames = int(cap.get(cv.CAP_PROP_FRAME_COUNT)) #获取帧数
        w = int(cap.get(3)) # 获取宽度 
        h = int(cap.get(4)) # 获取高度
         
        if num_frames < min_length:continue
        # if w!=320 or h!=240: continue
        
        if (cnt<num_train):
            dst_clips_path = os.path.join(data_root,'data','train',video_name)

        elif (num_train<=cnt<(num_train+num_val)):
            dst_clips_path = os.path.join(data_root,'data','val',video_name)

        elif cnt>=(num_train+num_val):
            dst_clips_path = os.path.join(data_root,'data','test',video_name)
        
        print('Preprocessing:', cnt, video_path)
        cnt = cnt + 1
        
        if cap.isOpened():
            frame_list = []  #set a list for img
            ret = True
            for i in range(num_frames):
                ret, frame = cap.read() #读取一帧图像
                if ret == True:frame_list.append(frame) #把帧放list
                else: break
            
            #创建每个视频的文件夹
            if not os.path.exists(dst_clips_path):
                os.mkdir(dst_clips_path)
            # 保存视频帧
            for i in range(len(frame_list)):
                frame_path = os.path.join(dst_clips_path, "%d.jpg"%i)  #get frame path
                # print(frame_path)
                cv.imwrite(frame_path, frame_list[i])             
    cap.release()    

split = (0.7,0.1,0.2) # proportion of train:valid:test datasets
min_length = 64
dataset_root = "/home/yanlq/zjut_mindvideo/mindvideo/datasets/UCF101"
raw_data_path = os.path.join(dataset_root, "raw_data")
label_save_path = os.path.join(dataset_root, "labels.txt")

create_label(raw_data_path, label_save_path)
# labels = pd.read_csv(os.path.join(dataset_root, 'labels.txt'))
# create_split_dir(dataset_root, labels)
prepare_clips(data_root=dataset_root, 
              raw_data_path=raw_data_path, 
              split=split, 
              min_length=min_length)